
#include "arg_max_with_value_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <iostream>

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  ArgMaxWithValueTilingData tiling;
    // auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    //auto socVersion = ascendcPlatform.GetSocVersion(); //获取当前硬件平台版本型号

    //ub 196352 
    //Cube核 20 
    //Vector核数 40
    // uint64_t ub_size;
    // ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size); //获取ub大小
    // uint32_t aivNum = ascendcPlatform.GetCoreNum(); //获取Vector内核数
    // ascendcPlatform.GetCoreNumAic() //AI Core上的Cube核 
    // ascendcPlatform.GetCoreNumAiv() //AI Core中Vector核数
    //std::cout << aivNum << " " << ub_size << " " << ascendcPlatform.GetCoreNumAic() << " " << ascendcPlatform.GetCoreNumAiv() << std::endl;

    // uint32_t inputBytes = GetSizeByDataType(context->GetInputDesc(0)->GetDataType()); //输入类型占用空间
    auto dt = context->GetInputDesc(0)->GetDataType();//输入类型

    auto x_shape = context->GetInputShape(0)->GetStorageShape();//获取运行时Shape
    auto x_shape_DimNum = x_shape.GetDimNum();
    int32_t x_shape_Dim[10];
    for(int i = 0; i < x_shape_DimNum; i++)
    {
        x_shape_Dim[i] = x_shape.GetDim(i);
    }

    int dimension_attr = *context->GetAttrs()->GetInt(0);
    // bool keep_dims_attr = *context->GetAttrs()->GetBool(1);

    int dimension = dimension_attr;
    if(dimension < 0) dimension += x_shape_DimNum;

    uint32_t num_cores = 1; //使用的内核数

    uint16_t compress[3] = {1,1,1};
    for(int i = 0; i < dimension; i++)
    {
        compress[0] *= x_shape_Dim[i];
    }
    compress[1] = x_shape_Dim[dimension];
    for(int i = dimension+1; i < x_shape_DimNum; i++)
    {
        compress[2] *= x_shape_Dim[i];
    }
    // std::cout << compress[0] << " " << compress[1] << " " << compress[2] << " " << dimension << std::endl;

    if(dt == ge::DT_FLOAT && compress[2] == 1 && compress[0] > 500)
    {
        // int32_t formerNum,tailLength;
        context->SetTilingKey(1);
        // if(compress[0] > 40)
        // {
            // formerNum = compress[0] % 40;
            // tailLength = compress[0] / 40;
            // num_cores = 40;
        // }
        // else
        // {
        //     formerNum = 0;
        //     tailLength = 1;
        //     num_cores = compress[0];
        // }
        compress[2] = compress[0] / 40;
        compress[0] = compress[0] % 40;
        num_cores = 40;
        // tiling.set_formerNum(formerNum);
        // tiling.set_tailLength(tailLength);
    }
    else
    {
        context->SetTilingKey(2);
        num_cores = 1;
    }
    tiling.set_compress0(compress[0]);
    tiling.set_compress1(compress[1]);
    tiling.set_compress2(compress[2]);

  context->SetBlockDim(num_cores); //设置参与计算的VectorCore或者CubeCore核数
  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
  return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class ArgMaxWithValue : public OpDef {
public:
    explicit ArgMaxWithValue(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_UINT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("indice")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("values")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_UINT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("dimension").Int();
        this->Attr("keep_dims").AttrType(OPTIONAL).Bool(false);

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(ArgMaxWithValue);
}
